#!/usr/bin/env python3
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################
"""Merged projects created by from-scratch and OFG core."""

import argparse
import json
import logging
import os
import shutil
from typing import Optional

logger = logging.getLogger(name=__name__)
LOG_FMT = ('%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] '
           ': %(funcName)s: %(message)s')


def retrieve_top_harness_index(benchmark_dir: str) -> Optional[str]:
  """Return the top harness index for a generated target"""
  result = []
  for path in os.listdir(os.path.join(benchmark_dir, 'status')):
    base_path = os.path.join(benchmark_dir, 'status', path)
    result_path = os.path.join(base_path, 'result.json')
    if os.path.isdir(base_path) and os.path.isfile(result_path):
      with open(result_path, 'r') as f:
        json_dict = json.load(f)
        if json_dict.get('compiles', False):
          result.append({
              'index': path,
              'coverage': json_dict.get('coverage', 0.0)
          })

  if result:
    return sorted(result, key=lambda item: (item.get('coverage')),
                  reverse=True)[0].get('index')

  return None


def retrieve_success_harness_name(project_dir_list: list[str]) -> list[str]:
  """Returns name of success harnesses of the target project."""
  result_list = []
  for project_dir in project_dir_list:
    top_harness_index = retrieve_top_harness_index(project_dir)
    if top_harness_index:
      result_list.append(f'{os.path.basename(project_dir)}-{top_harness_index}')

  return result_list


def group_result_dir_by_project_name(
    result_dir: str, project_name_list: list[str]) -> dict[str, list]:
  """This function group the generater directory by project name."""
  result_dir_map = {}
  for project_basename in os.listdir(result_dir):
    for project_name in project_name_list:
      if project_basename.startswith(f'output-{project_name}-'):
        project_list = result_dir_map.get(project_name, [])
        project_list.append(os.path.join(result_dir, project_basename))
        result_dir_map[project_name] = project_list
        break

  return result_dir_map


def copy_success_ofg_autogens(result_dir: str, destination: str,
                              project_name_list: list[str],
                              oss_fuzz_path: str) -> None:
  """Copies the success harnesses for each projects to destination."""
  result_dir_map = group_result_dir_by_project_name(os.path.abspath(result_dir),
                                                    project_name_list)

  for project_name, project_dir_list in result_dir_map.items():
    logger.info('Handling project %s', project_name)
    destination_dir = os.path.join(destination, project_name)
    if not os.path.isdir(destination_dir):
      os.mkdir(destination_dir)

    top_harnesses = retrieve_success_harness_name(project_dir_list)
    logger.info('Found %d success harnesses for project %s', len(top_harnesses),
                project_name)
    for name in top_harnesses:
      src_dir = os.path.join(oss_fuzz_path, 'projects', name)
      dst_dir = os.path.join(destination_dir, name.rsplit('-', 1)[0])
      shutil.copytree(src_dir, dst_dir)

    if len(os.listdir(destination_dir)) == 0:
      logger.warning('No success target for project %s', project_name)
      shutil.rmtree(destination_dir)


def parse_commandline():
  """Parse commandline."""
  parser = argparse.ArgumentParser()
  parser.add_argument('--result-dir',
                      '-r',
                      help='Results created by OFG.',
                      type=str)
  parser.add_argument(
      '--destination-dir',
      '-d',
      help='Folder with projects generated by from-scratch OFG.',
      type=str)
  parser.add_argument(
      '--project-name',
      '-p',
      help='A comma separated string of all target project name.',
      type=str)
  parser.add_argument('--oss-fuzz-path',
                      '-o',
                      help='Path of the OSS-Fuzz used by OFG.',
                      type=str)
  return parser.parse_args()


def main():
  """CLI entrypoint."""
  logging.basicConfig(level=logging.INFO, format=LOG_FMT)

  args = parse_commandline()
  copy_success_ofg_autogens(args.result_dir, args.destination_dir,
                            args.project_name, args.oss_fuzz_path)


if __name__ == "__main__":
  main()
